import base64
import os
import time
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime

from PIL import Image
from flask import g
import io
from flask import jsonify
from werkzeug.utils import secure_filename

from DBService.Working_Service import WorkingService
from routes.app import app
from flask import request, send_file

from routes.segment_routes import allowed_file
from service.segment.runSeg import generate_image_number
from service.styleCLIP.PictureInitialization import process_image
from service.styleCLIP.run_latent_optimization import getArgs, runStyleCLIP, translate_description


# 对图片进行初始化
@app.route('/styleCLIP/PictureInitalization', methods=['POST'])
def PictureInitalization():
    # 从请求中获取image_path, latent_save_path参数
    image_path = request.args.get('image_path')
    print(image_path)
    latent = 'C:\\Users\\Dell\\PycharmProjects\\ThinkStation\\service\\styleCLIP\\Parameter'
    latent_number = generate_image_number(latent)
    latent_save_path = os.path.join(latent, f'{latent_number}')
    print(latent_save_path)
    # 调用process_image函数对图片进行向量初始化(代码本身设计成懒加载模式，所以不需要初始化)
    latent_save_path,image_name = process_image(image_path, latent_save_path)
    # 返回初始化完成的消息和latent_save_path以及image_name
    return jsonify({"message": "图片初始化完成", "latent_save_path": latent_save_path, "image_name": image_name})
    #return jsonify({'message': '图片初始化完成', 'latent_save_path': latent_save_path})




# 根据初始化后的潜在下潜在向量来对图片的属性进行篡改
@app.route('/styleCLIP/run_latent_optimization', methods=['POST'])
def run_latent_optimization():
    # 获取参数
    data = request.get_json()
    latent_save_path = data.get('latent_save_path',
                                'C:\\Users\\Dell\\PycharmProjects\\ThinkStation\\service\\styleCLIP\\Parameter\\latent_code669.pt')
    print(latent_save_path)
    image_name = data.get('image_name')
    print(image_name)
    description = data.get('description', 'a person with purple hair')
    description_cn = description
    stylegan_size = data.get('stylegan_size', 1024)
    lr_rampup = data.get('lr_rampup', 0.05)
    lr = data.get('lr', 0.1)
    optimization_steps = data.get('optimization_steps', 3)
    experiment_type = data.get('experiment_type', 'edit')
    l2_lambda = data.get('l2_lambda', 0.008)
    id_lambda = data.get('id_lambda', 0.006)
    stylespace = data.get('stylespace', False)
    truncation = data.get('truncation', 0.7)
    create_video = data.get('create_video', False)

    # 获取参数
    argsn = getArgs(description_cn, stylegan_size, lr_rampup, lr,
                   optimization_steps, experiment_type, l2_lambda,
                   id_lambda, stylespace, latent_save_path, truncation,
                   create_video)

    # 运行StyleCLIP
    result = runStyleCLIP(image_name,**argsn)

    # 将PIL Image对象转换为字节流
    byte_arr = io.BytesIO()
    result.save(byte_arr, format='PNG')
    byte_arr = byte_arr.getvalue()

    # 将字节流编码为base64
    result_encoded = base64.b64encode(byte_arr).decode('utf-8')
    # 返回图片
    return jsonify({"message": "篡改完成", "result_encoded": result_encoded, "image_name": image_name})
# 调用getArgs函数获取参数
    # description_cn: 描述，用于生成图片的文本描述，这里是中文描述的英文翻译
    # stylegan_size: StyleGAN2模型的生成图片的大小
    # lr_rampup: 学习率上升期，即在这个步骤数内，学习率从0线性增长到lr
    # lr: 学习率
    # optimization_steps: 优化步骤数
    # experiment_type: 实验类型，可以是"edit"或"free_generation"
    # l2_lambda: L2正则化的权重
    # id_lambda: 身份损失的权重
    # stylespace: 是否在样式空间中工作，如果为False，则在潜在空间中工作
    # latent_save_path: 潜在向量的保存路径
    # truncation: 截断值，用于控制生成图片的多样性和质量的平衡
    # create_video: 是否创建视频，如果为True，则每一步都保存图片，否则每20步保存一次图片


#如果优化后的图片满足要求，则执行保存图片操作
@app.route('/styleCLIP/save_image', methods=['POST'])
def save_image():
    print("Files:", request.files)
    print("Form data:", request.form)
    file = request.files.get('File')
    if file:
        # 获取文件的原始名字和扩展名
        filename, file_extension = os.path.splitext(secure_filename(file.filename))
        # 添加时间戳到文件名
        filename = f"{filename}_{int(time.time())}{file_extension}"
        print(filename)
        # 检查文件名是否为空
        if file.filename == '':
            return jsonify({"error": "没有选择文件"}), 402

        # 检查文件类型
        if not allowed_file(file.filename):
            return jsonify({"error": "不允许的文件类型"}), 403
        # 保存文件
        dir_path = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'image'))
        print(dir_path)
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        save_path = os.path.join(dir_path, filename)
        try:
            file.save(save_path)
            print('图片已保存' + save_path)
            file_size = os.path.getsize(save_path)
            uid = request.form.get('uid')
            print(uid)
            # 将图片信息存入数据库
            relative_save_path = os.path.relpath(save_path, start=os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
            image=WorkingService().add_DB_image(uid, relative_save_path,filename)
            print('2')
            if image is not None:
                return jsonify(
                    {"message": "文件上传成功", "filename": filename, "size": file_size}), 200  # 使用 jsonify 返回 JSON 响应
            else:
                return jsonify({"error": "文件上传失败"}), 500
        except Exception as e:

            return jsonify({"error": "保存文件或存入数据库时发生错误: " + str(e)}), 500
        else:
            return jsonify({"message": "请求中没有文件部分."}), 400



@app.route('/styleCLIP/save_image2', methods=['POST'])
def save_image2():
    print("Files:", request.files)
    print("Form data:", request.form)
    file = request.files.get('File')
    if file:
        # 获取文件的原始名字和扩展名
        filename, file_extension = os.path.splitext(secure_filename(file.filename))
        # 添加时间戳到文件名
        filename = f"{filename}_{int(time.time())}{file_extension}"
        print("filename")
        print(filename)
        # 检查文件名是否为空
        if file.filename == '':
            return jsonify({"error": "没有选择文件"}), 402

        # 检查文件类型
        if not allowed_file(file.filename):
            return jsonify({"error": "不允许的文件类型"}), 403
        # 保存文件
        dir_path = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'image'))
        print(dir_path)
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        save_path = os.path.join(dir_path, filename)
        try:
            file.save(save_path)
            print('图片已保存' + save_path)
            # 获取用户ID
            uid = request.form.get('uid')

            # 获取Watermark参数
            watermark = request.form.get('Watermark')
            print("Watermark:")
            print(watermark)

            # 如果Watermark为1，则添加水印
            if watermark == '1':
                # 加载图片
                img = Image.open(save_path)

                # 创建一个ImageDraw对象
                draw = ImageDraw.Draw(img)
                # 创建水印文本
                watermark_text = f"{uid}_{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"

                # 创建字体对象，第一个参数是字体文件的路径，第二个参数是字体大小
                font = ImageFont.truetype('C:\\Windows\\Fonts\\arial.ttf', 15)

                # 在图片的左上角添加水印
                draw.text((0, 0), watermark_text, fill="black", font=font)

                # 保存带有水印的图片
                img.save(save_path)

            file_size = os.path.getsize(save_path)
            print("uid")
            print(uid)
            # 将图片信息存入数据库
            relative_save_path = os.path.relpath(save_path, start=os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
            image=WorkingService().add_DB_image(uid, relative_save_path,filename)
            print('2')
            if image is not None:
                return jsonify(
                    {"message": "文件上传成功", "filename": filename, "size": file_size}), 200  # 使用 jsonify 返回 JSON 响应
            else:
                return jsonify({"error": "文件上传失败"}), 500
        except Exception as e:

            return jsonify({"error": "保存文件或存入数据库时发生错误: " + str(e)}), 500
        else:
            return jsonify({"message": "请求中没有文件部分."}), 400
